import cv2

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1024)
# cap.set(cv2.CAP_PROP_FPS, 30)
# cap.set(cv2.CAP_PROP_BRIGHTNESS, 10)
# cap.set(cv2.CAP_PROP_CONTRAST, 40)
# cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0.25)
# cap.set(cv2.CAP_PROP_EXPOSURE, 50)

CHESS_WIDTH = 11
CHESS_HEIGHT = 8

def detect_chessboard( image):
    """检测棋盘格"""
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    ret, corners = cv2.findChessboardCorners(gray, (CHESS_WIDTH, CHESS_HEIGHT), None)
    if ret:
        # 角点精检测
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        # Draw and display the corners
        image = cv2.drawChessboardCorners(image, (CHESS_WIDTH, CHESS_HEIGHT), corners2, ret)
    return image
# cap2 = cv2.VideoCapture(1)

def main(output_dir=''):
    i = 0
    while(cap.isOpened()):
        ret, frame = cap.read()
        image = detect_chessboard(frame.copy())
        cv2.imshow('frame', image)
        # ret, frame2 = cap2.read()
        # print("frame2", frame2.shape)
        # cv2.imshow('frame2', frame2)

        i += 1
        print(i)
        key = cv2.waitKey(100000)
        if key == ord('n'):
            continue
        if key == ord('q'):
            break
        if key == ord('s'):
            cv2.imwrite(f'{output_dir}/frame_{i}.jpg', frame)
    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    # remove all files in images_tmp
    import os
    if os.path.exists("images_tmp"):
        for file in os.listdir("images_tmp"):
            os.remove(f"images_tmp/{file}")
    else:
        os.makedirs("images_tmp")
    main(output_dir="images_tmp")
